perf(gated_delta_net): fold q/k L2-norm into the gated_delta_rule kernel - #5396
perf(gated_delta_net): fold q/k L2-norm into the gated_delta_rule kernel#5396yuchenwang3 wants to merge 2 commits into
Conversation
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
|
@xuantengh pinging you as the most recent committer on |
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
|
Thanks for the contribution! We had a refactor to the GDN module recently in #5843, would you please rebase that and apply this to the new file? |
svcnvidia-nemo-ci
added
the
waiting-on-customer
|
Done, @xuantengh — merged latest
Net diff vs main is the same two-line semantic change as before (+7/−4 with the comment block). Ready for another look. |
svcnvidia-nemo-ci
removed
the
waiting-on-customer
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
|
@xuantengh one more follow-up before I refresh this again: current main added GDN2 sharing |
b752055 to
39fd9a0
Compare
|
Rebased onto current main and adjusted the shared input helper added since the original PR: GDN now defers q/k normalization to its kernel, while GDN2 keeps the existing caller-side normalization. Local formatting, syntax, and routing checks pass. @xuantengh, could you re-review this head when you get a chance? |
svcnvidia-nemo-ci
removed
the
waiting-on-maintainers
|
Sorry for the delayed reply. Would you please add a test to |
Defer GDN q/k normalization to the gated delta rule kernel so the kernel owns normalization and caller autograd no longer retains an additional pre-split normalized query_key activation alongside the kernel q/k. Keep caller-side normalization as the default for GDN2, which shares the input-preparation helper but does not enable this path. Signed-off-by: Yuchen Wang <yw.yy953e@alibaba-inc.com>
39fd9a0 to
8206dc3
Compare
FLA l2norm normalizes the last dimension and does not accept a dim argument. Use its default eps=1e-6 API in the deterministic GDN and GDN2 fallbacks. Add a BF16 end-to-end regression that compares the former caller-normalized route with kernel-owned normalization across outputs, input gradients, and parameter gradients, while asserting the helper and kernel routing flags. Signed-off-by: Yuchen Wang <yw.yy953e@alibaba-inc.com>
8206dc3 to
398a42b
Compare
|
Done, @xuantengh. I added a small kernel-level parity check plus a BF16 end-to-end test that runs the real The new tests pass on all 8 ranks on 8×B200. While running the existing deterministic matrix I also caught and fixed a current FLA API mismatch in the newly activated fallback ( |
4 participants
What
Move GatedDeltaNet q/k L2 normalization into the FLA
gated_delta_rulekernel instead of materializing normalized q/k before the call.The current GDN and GDN2 implementations share
_prepare_input_for_gated_delta_rule, but GDN2's kernel path does not enable in-kernel normalization. The rebased change therefore makes deferral an explicit helper option:use_qk_l2norm_in_kernel=self.use_qk_l2norm.Falseto its kernel.This preserves GDN2 behavior added after the original PR rather than silently removing its normalization during the rebase.
Why
Caller-side normalization leaves an additional pre-split normalized
query_keyactivation live for backward alongside the q/k tensors the gated-delta kernel already needs. Letting FLA own normalization removes that extra caller-side activation and keeps the normalization state with the kernel-owned q/k, which saves activation memory in long-context GDN-heavy hybrids.The numerical contract is unchanged:
l2normand the in-kernel path useeps=1e-6;use_qk_l2norm=Falseremains a no-op.Validation
The original change ran in Qwen3.5-35B-A3B 128K SFT on 16×B200. The current-main replay passes Black, isort, Ruff, Pylint,
py_compile, andgit diff --check.The final
398a42b4head was also exercised with an 8-rank torchrun on 8×B200:GatedDeltaNet.forwardcomparison passed on every rank, covering outputs, input gradients, all parameter gradients, and both routing flags;